# -*- coding: utf-8 -*-
import numpy as np
import os
import wrf_output
import gauge
import matplotlib.pyplot as plt


def decumulation(seq):
    result = [seq[i + 1] - seq[i] for i in range(len(seq) - 1)]
    result.insert(0, seq[0])
    return result


def enlarge_xy_label():
    for x_label in plt.gca().xaxis.get_ticklabels():
        x_label.set_fontsize(20)
    for y_label in plt.gca().yaxis.get_ticklabels():
        y_label.set_fontsize(20)


def get_statistic(dat_path, domain, start_time, end_time):
    # get every mp_cu group wrf_output and extract the center value
    mps = ['2', '6', '7']
    cus = ['0', '1', '2', '7']
    pbls = ['2', '6', '9']
    split_path = dat_path.split('/')
    gridSpacing = split_path[-1]
    gauge_loc = ''
    if gridSpacing == '139':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge' \
                    r'/Gauge_row_col_1.csv'
    elif gridSpacing == '51545':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge' \
                    r'/Gauge_row_col_5.csv'
    elif gridSpacing == '103090':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge' \
                    r'/Gauge_row_col_10.csv'
    simulations = []
    for i in mps:
        for j in cus:
            for k in pbls:
                mp_cu_pbl = i + j + k
                if os.path.exists(dat_path + '/' + mp_cu_pbl):
                    input_path = dat_path + '/' + mp_cu_pbl + \
                        '/wrfout_d' + domain + '_' + start_time
                    all_gauges = wrf_output.GetAll(input_path, gauge_loc)
                    all_gauges = decumulation(all_gauges)
                    simulations.append(all_gauges)
    simulations = np.asarray(simulations)
    min = np.min(simulations, axis=0)
    max = np.max(simulations, axis=0)
    band = np.array([min, max])
    band_50 = np.mean(band, axis=0)
    band_50 = band_50.astype(np.float64)
    # get the center gauge value
    gauge50 = gauge.Get50Gauge(start_time, end_time)

    # calculate dispersion statistic when ensemble mean larger than 0
    band_mask = (band_50 < 0)
    band_95 = np.sum(band, axis=0) * 0.95
    band_95 = band_95.astype(np.float64)
    band_5 = np.sum(band, axis=0) * 0.05
    band_5 = band_5.astype(np.float64)
    band_95_filtered = np.ma.array(band_95, mask=band_mask)
    band_5_filtered = np.ma.array(band_5, mask=band_mask)
    # dispersion statistic
    dispersion_d = np.mean((band_95_filtered - band_5_filtered), axis=0)
    # normalized dispersion statistic
    # denominator = band_50.copy()
    # denominator[denominator == 0] = np.nan
    # nd = (band_95 - band_5) / denominator
    # masked_nd = np.ma.masked_array(nd, np.isnan(nd))
    # dispersion_nd = np.mean(masked_nd, axis=0).filled(np.nan)
    # print result
    # print("Dispersion statistic:\n", dispersion_d)
    # print("Normalized dispersion statistic:\n", dispersion_nd)
    # calculate MAE and MSE
    # mae = np.abs(np.mean((band_50 - gauge50.values), axis=0))
    mae = np.mean(np.abs((band_50 - gauge50.values)), axis=0)
    mse = np.mean(np.power((band_50 - gauge50.values), 2), axis=0)

    # print("MAE:\n", mae)
    # print("MSE:\n", mse)
    # output statistic result
    all_statistic = np.array([dispersion_d, mae, mse])
    all_statistic = all_statistic.T
    return all_statistic
    # output_path = r'F:/Experiment/Research/Result/' + \
    #     split_path[3] + '/' + split_path[4] + '/' + \
    #     split_path[5] + '/' + split_path[6] + '/'
    # if not os.path.exists(output_path):
    #     os.makedirs(output_path)
    # np.savetxt(output_path + split_path[6] + '.csv', all_statistic,
    #            delimiter=',', fmt='%.2f', header='Dispersion,MAE,MSE')


def plot_box(dat_path, domain, start_time, end_time, event):
    scenarios = ['139', '51545', '103090']
    scenarios_stat = []
    for scenario in scenarios:
        input_path = dat_path + '/' + scenario
        stat = get_statistic(input_path, domain, start_time, end_time)
        scenarios_stat.append(stat)
    scenarios_stat = np.array(scenarios_stat)
    dispersion = scenarios_stat[:, :, 0].T
    mae = scenarios_stat[:, :, 1].T
    # mse = scenarios_stat[:, :, 2].T

    # plot box
    plt.style.use('fivethirtyeight')
    fig = plt.figure(figsize=(20, 10))
    fig.suptitle(event, fontsize=24, fontweight='bold')
    # plot dispersion statistic
    plt.subplot(1, 2, 1)
    plt.boxplot(dispersion, 0, showmeans=True, patch_artist=True,
                medianprops={'color': '#297083'},
                boxprops={'color': '#539caf', 'facecolor': '#539caf'},
                whiskerprops={'color': '#539caf'},
                labels=['1', '2', '3'])
    plt.title('Dispersion statistic', fontsize=22)
    # plt.xlabel('Scenarios', fontdict={'size': 22})
    plt.ylabel('Dispersion', fontdict={'size': 22})
    enlarge_xy_label()
    # plot MAE
    plt.subplot(1, 2, 2)
    plt.boxplot(mae, 0, showmeans=True, patch_artist=True,
                medianprops={'color': '#297083'},
                boxprops={'color': '#539caf', 'facecolor': '#539caf'},
                whiskerprops={'color': '#539caf'},
                labels=['1', '2', '3'])
    plt.title('MAE statistic', fontsize=22)
    # plt.xlabel('Scenarios', fontdict={'size': 22})
    plt.ylabel('MAE', fontdict={'size': 22})
    plt.text(0.3, -0.13, 'Scenarios', ha='center', fontdict={'size': 22})
    enlarge_xy_label()
    plt.show()


def main():
    """main function"""
    # new run
    # R1
    dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/01/1700'
    domain = r'03'
    start_time = r'2008-01-17_00_00_00'
    end_time = r'2008-01-19_12_00_00'
    event = 'R1'

    # R2
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/01/1912'
    # domain = r'03'
    # start_time = r'2008-01-19_12_00_00'
    # end_time = r'2008-01-22_00_00_00'
    # event = 'R2'

    # R3
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/08/1718'
    # domain = r'03'
    # start_time = r'2008-08-17_18_00_00'
    # end_time = r'2008-08-20_00_00_00'
    # event = 'R3'

    # R4
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/09/0500'
    # domain = r'03'
    # start_time = r'2008-09-05_00_00_00'
    # end_time = r'2008-09-07_00_00_00'
    # event = 'R4'

    # R5
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/09/2900'
    # domain = r'03'
    # start_time = r'2008-09-29_00_00_00'
    # end_time = r'2008-10-02_06_00_00'
    # event = 'R5'

    # R6
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/10/2506'
    # domain = r'03'
    # start_time = r'2008-10-25_06_00_00'
    # end_time = r'2008-10-26_06_00_00'
    # event = 'R6'

    # R7
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/11/0900'
    # domain = r'03'
    # start_time = r'2008-11-09_00_00_00'
    # end_time = r'2008-11-10_06_00_00'
    # event = 'R7'

    # R8
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/12/0400'
    # domain = r'03'
    # start_time = r'2008-12-04_00_00_00'
    # end_time = r'2008-12-06_00_00_00'
    # event = 'R8'

    plot_box(dat_path, domain, start_time, end_time, event)


if __name__ == '__main__':
    main()
